use hashbrown::HashMap;

fn main() {
    // load onnx model
    let bytes = std::fs::read("example/const_chain_f32.onnx").unwrap();
    let model = zhulong::onnx::load(&bytes).unwrap();
    let opset = zhulong::onnx::opset_ai_onnx(&model).unwrap();

    println!("opset: {opset}");

    // parse to IR
    let mut graph = zhulong::ir::graph::Graph::from_onnx(&model);
    println!(
        "{}",
        graph.pretty(zhulong::ir::graph::PrettyOpts {
            show_attrs: true,
            show_constants: true,
            const_max_elems: 16,
        })
    );

    // run optimization passes
    let changed = zhulong::pass::pipeline::run_all_passes(&mut graph);
    if changed {
        println!("after optimization passes:");
        println!(
            "{}",
            graph.pretty(zhulong::ir::graph::PrettyOpts {
                show_attrs: true,
                show_constants: true,
                const_max_elems: 16,
            })
        );
    }

    //
    let ops = zhulong::ops::register::register_all_ops();
    println!("registered ops: {:?}", ops);

    // build input tensors
    let x = zhulong::ir::tensor::Tensor::from_f32(&[2, 2], vec![1., 2., 3., 4.]);
    let y = zhulong::ir::tensor::Tensor::from_f32(&[2, 2], vec![10., 20., 30., 40.]);

    println!("x = {:?}", x.as_slice::<f32>());
    println!("y = {:?}", y.as_slice::<f32>());

    //
    let mut feeds: HashMap<String, zhulong::ir::tensor::Tensor> = HashMap::new();
    feeds.insert("x".into(), x);
    feeds.insert("y".into(), y);

    //
    let outs = zhulong::runtime::executor::run(&graph, feeds);
    let z = outs.get("z").unwrap();

    println!("z = {:?}", z.as_slice::<f32>());
}
